import numpy as np
from matplotlib import pyplot as plt


def forward(w, x):
    return w * x


def loss(x_data, y_data, w):
    cost = 0
    for x, y in zip(x_data, y_data):
        y_pred = forward(w, x)
        cost += (y_pred - y) * (y_pred - y)
    return cost / len(x_data)


def grad(x_data, y_data, w):
    w_grad = 0
    for x, y in zip(x_data, y_data):
        w_grad += 2 * (y - forward(w, x)) * (-x)
    return w_grad/len(x_data)


if __name__ == '__main__':
    x_data = [1.0, 2.0, 3.0]
    y_data = [2.0, 4.0, 6.0]

    epoch = 100
    lr = 0.01
    w = 1
    loss_list = []
    for i in np.arange(epoch):
        loss_val = loss(x_data, y_data, w)
        w -= lr * grad(x_data, y_data, w)
        loss_list.append(loss_val)
        print("epoch=", i, "w=", w, "loss=", loss_val)

    plt.figure()
    plt.xlabel("epoch")
    plt.ylabel("loss")
    plt.plot(loss_list)
    plt.show()
